In [1]:
import sympy as sp
import numpy as np
%matplotlib nbagg
from matplotlib import pyplot as plt
class hamiltonian_system:
def __init__(self, x, p, H):
assert(len(x) == len(p))
self.H = H
self.x = x
self.p = p
self.degrees_of_freedom = len(self.x)
self.rhs = ([sp.utilities.lambdify(tuple(self.x)+tuple(self.p),
self.H.diff(self.p[i]),
modules = 'numpy')
for i in range(self.degrees_of_freedom)] +
[sp.utilities.lambdify(tuple(self.x)+tuple(self.p),
-self.H.diff(self.x[i]),
modules = 'numpy')
for i in range(self.degrees_of_freedom)])
self.energy = sp.utilities.lambdify(tuple(self.x)+tuple(self.p),
self.H,
modules = 'numpy')
return None
def numpy_rhs(self, point):
return np.array([self.rhs[i](*tuple(point)) * np.ones(point.shape[1:])
for i in range(2*self.degrees_of_freedom)])
In [2]:
def Euler(rhs, dt = 0.5, N = 8, x0 = None):
x = np.zeros((N+1,) + x0.shape,
dtype = x0.dtype)
x[0] = x0
for t in range(N):
x[t+1] = x[t] + dt*rhs(x[t])
return x
def cRK(rhs, dt = 0.5, N = 8, x0 = None):
x = np.zeros((N+1,) + x0.shape,
dtype = x0.dtype)
x[0] = x0
for t in range(N):
k1 = rhs(x[t])
k2 = rhs(x[t] + 0.5*dt*k1)
k3 = rhs(x[t] + 0.5*dt*k2)
k4 = rhs(x[t] + dt*k3)
x[t+1] = x[t] + dt*(k1 + 2*(k2 + k3) + k4)/6
return x
def get_evdt(
initial_condition,
h0, N0, ndivisions,
method,
rhs):
epsilon = []
tstep = [h0]
sol = []
N = N0
sol.append(method(rhs,
dt = tstep[-1],
N = N,
x0 = initial_condition))
for n in range(1, ndivisions + 1):
tstep.append(h0*2.**(-n))
N = N*2
sol.append(method(rhs,
dt = tstep[-1],
N = N,
x0 = initial_condition))
epsilon.append((sol[-1][-1] - sol[-2][-1]) /
np.max(np.abs(sol[-1])))
return (np.array(tstep[:len(tstep)-1]),
np.abs(epsilon))
In [3]:
nparticles = 8
position = []
momentum = []
for i in range(nparticles):
position += [sp.Symbol('x{0}'.format(i)),
sp.Symbol('y{0}'.format(i)),
sp.Symbol('z{0}'.format(i))]
momentum += [sp.Symbol('p{0}'.format(i)),
sp.Symbol('q{0}'.format(i)),
sp.Symbol('r{0}'.format(i))]
def VanDerWaals(r1, r2):
dist = sp.sqrt((r1[0] - r2[0])**2 + (r1[1] - r2[1])**2 + (r1[2] - r2[2])**2)
return dist**(-12) - dist**(-6)
def Box(r, exponent = 4):
return ((r[0]/(nparticles))**(2*exponent) +
(r[1]/(nparticles))**(2*exponent) +
(r[2]/(nparticles))**(2*exponent))
# first, put in kinetic energy term
H = sum(p**2/2 for p in momentum)
# second, put in walls
H += sum(Box(position[3*i:3*(i+1)]) for i in range(nparticles))
# third, put in pairwise interactions
for i in range(nparticles):
for j in range(i+1, nparticles):
H += VanDerWaals(position[3*i:3*(i+1)],
position[3*j:3*(j+1)])
bla = hamiltonian_system(position, momentum, H)
In [4]:
initial_condition = np.zeros(nparticles*6, dtype = np.float)
initial_condition[0:3*nparticles:3] = \
np.linspace(-nparticles/2, nparticles/2, nparticles)
alpha = np.random.random()*np.pi*2
beta = np.random.random()*np.pi
initial_condition[3*nparticles ] = 2*(np.sin(beta)*np.cos(alpha)+2)
initial_condition[3*nparticles+1] = 2*(np.sin(beta)*np.sin(alpha))
initial_condition[3*nparticles+2] = np.cos(beta)
print initial_condition
In [5]:
x = np.linspace(0.9, 3, 128)
a = plt.figure(figsize = (6, 3)).add_subplot(111)
a.plot(x, x**(-12) - x**(-6))
a.plot(x, x - x, color = (0, 0, 0), dashes = (2, 2))
a.set_title('Van der Waals potential')
# basically, particles will attract if they're further apart than 1.1,
# and repell otherwise. in practice, the attraction is negligible if the distance
# is larger than 2 (since the derivative of U is teeny).
Out[5]:
In [6]:
dt, err = get_evdt(initial_condition,
2.**(-6), 16, 6,
cRK,
bla.numpy_rhs)
fig = plt.figure()
a = fig.add_subplot(111)
a.plot(dt, np.average(err, axis = 1), marker = '.')
a.set_xscale('log')
a.set_yscale('log')
In [7]:
from mayavi import mlab
sol = cRK(bla.numpy_rhs,
dt = 2.**(-6),
N = 2**7,
x0 = initial_condition)
In [8]:
points = mlab.points3d(sol[0, 0:3*nparticles:3],
sol[0, 1:3*nparticles:3],
sol[0, 2:3*nparticles:3],
scale_mode = 'none',
scale_factor = 0.9)
# words written with "@" in front are called decorators,
# and they are basically functions that are applied to the function
# defined underneath, as in https://www.python.org/dev/peps/pep-0318/#id71
# something like that.
# if you have the patience, you can read the whole thing
# https://www.python.org/dev/peps/pep-0318
# http://docs.enthought.com/mayavi/mayavi/mlab_animating.html
@mlab.show
@mlab.animate(delay=250)
def anim():
t = 0
while 1:
points.mlab_source.set(x = sol[t, 0:3*nparticles:3],
y = sol[t, 1:3*nparticles:3],
z = sol[t, 2:3*nparticles:3])
t = (t + 1) % sol.shape[0]
yield
anim()
Very nice explanation of the yield keyword is available at
http://stackoverflow.com/a/231855/4205267.
Basically, anim returns a "generator", a function that will act like a list that can only be traversed once (conceptually a queue).
What happens in practice is that the while loop is executed once for each frame, and that's what we care about.
In [19]:
initial_condition = np.zeros(nparticles*6, dtype = np.float)
initial_condition[3*nparticles:] = np.random.randn(nparticles*3)
initial_condition[0:3*nparticles:3] = \
np.linspace(-nparticles/2, nparticles/2, nparticles)
sol = cRK(bla.numpy_rhs,
dt = 2.**(-7),
N = 2**12,
x0 = initial_condition)
In [20]:
points = mlab.points3d(sol[0, 0:3*nparticles:3],
sol[0, 1:3*nparticles:3],
sol[0, 2:3*nparticles:3],
scale_mode = 'none',
scale_factor = 0.9)
@mlab.show
@mlab.animate(delay=50)
def anim():
t = 0
while 1:
points.mlab_source.set(x = sol[t, 0:3*nparticles:3],
y = sol[t, 1:3*nparticles:3],
z = sol[t, 2:3*nparticles:3])
t = (t + 10) % sol.shape[0]
yield
anim()
In [ ]: